package aqm

import (
	"context"

	"abm-cache-center/rate"
	bm "abm-cache-center/net/blade"
	"abm-cache-center/net/blade/ecode"
)

// AQM aqm midleware.
type AQM struct {
	limiter rate.Limiter
}

// New new and return supervisor midleware.
func New(limiter rate.Limiter) (s *AQM) {
	s = &AQM{
		limiter: limiter,
	}
	return
}

func (a *AQM) ServeHTTP(c *bm.Context) {
	done, err := a.limiter.Allow(c)
	if err != nil {
		// TODO: priority request.
		c.JSON(nil, err)
		return
	}
	defer func() {
		if c.Error != nil && !ecode.Deadline.Equal(c.Error) && c.Err() != context.DeadlineExceeded {
			done(rate.Ignore)
			return
		}
		done(rate.Success)
	}()
	c.Next()
}
